from flask import Flask, request, jsonify
from flask_cors import CORS
import joblib
import traceback
import os

# 构建指向当前文件所在目录的绝对路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# 使用绝对路径构建模型文件的完整路径
ABBE_MODEL_PATH = os.path.join(BASE_DIR, "abbe_predictor_20241216")
RI_MODEL_PATH = os.path.join(BASE_DIR, "RI_predictor_20241216")

# 从你的模型文件中导入函数
from model import ringed_SMILES_gen, labelling_optical

app = Flask(__name__)
CORS(app) 

# 加载机器学习模型
try:
    abbe_predictor = joblib.load(ABBE_MODEL_PATH)
    print(f"OK: Abbe number predictor loaded successfully from: {ABBE_MODEL_PATH}")
except Exception:
    abbe_predictor = None
    print(f"WARNING: Abbe predictor not found or failed to load from '{ABBE_MODEL_PATH}'.")

try:
    ri_predictor = joblib.load(RI_MODEL_PATH)
    print(f"OK: Refractive index predictor loaded successfully from: {RI_MODEL_PATH}")
except Exception:
    ri_predictor = None
    print(f"WARNING: RI predictor not found or failed to load from '{RI_MODEL_PATH}'.")


# API端点，已支持多组分输入
@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json()
    if not data or 'molecules' not in data or 'properties' not in data:
        return jsonify({'error': 'Invalid input. "molecules" and "properties" are required.'}), 400

    molecules = data['molecules']
    props_to_predict = data['properties']
    
    if not isinstance(molecules, list) or len(molecules) == 0:
        return jsonify({'error': 'Input "molecules" must be a non-empty list.'}), 400
    
    for mol in molecules:
        if not all(k in mol for k in ('smiles', 'ratio')):
            return jsonify({'error': 'Each item in "molecules" must contain "smiles" and "ratio".'}), 400
        if not mol['smiles']:
            return jsonify({'error': 'SMILES string cannot be empty for any component.'}), 400

    results = {}
    
    try:
        if len(molecules) == 1:
            smiles = molecules[0]['smiles']
            processed_smiles = smiles
            if '*' in smiles:
                processed_smiles, _ = ringed_SMILES_gen(smiles)
            results['processed_smiles'] = processed_smiles
            feature_smiles = processed_smiles
        else:
            component_strings = [f"{mol['smiles']} ({mol['ratio']:.1f}%)" for mol in molecules]
            results['mixture_info'] = " / ".join(component_strings)
            # 演示用：对第一个组分进行预测
            first_smiles = molecules[0]['smiles']
            if '*' in first_smiles:
                feature_smiles, _ = ringed_SMILES_gen(first_smiles)
            else:
                feature_smiles = first_smiles

        # 属性预测循环
        for prop in props_to_predict:
            if prop == "阿贝数（演示用）":
                if abbe_predictor and 'feature_smiles' in locals():
                    feature = labelling_optical(feature_smiles, 1)
                    abbe = abbe_predictor.predict([feature])
                    results['阿贝数（演示用）'] = f"{abbe[0]:.4f}"
                else:
                    results['阿贝数（演示用）'] = "Error: Model not available."
            
            elif prop == "折射率":
                if ri_predictor and 'feature_smiles' in locals():
                    feature = labelling_optical(feature_smiles, 1)
                    ri = ri_predictor.predict([feature])
                    results['折射率'] = f"{ri[0]:.4f}"
                else:
                    results['折射率'] = "Error: Model not available."
            
            elif prop in ["均方回转半径", "LogP", "熔点", "热分解温度", "热膨胀系数", "密度", "粘度", "弹性模量"]:
                results[prop] = "模型开发中..."

        return jsonify(results)

    except Exception as e:
        print(traceback.format_exc())
        return jsonify({'error': f'An error occurred during prediction: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)